MUM 2024-25 Modele dyfuzyjne¶
%matplotlib inline
#Required for Jupyter to display plots as output
# inline -→> display inLine for Jupyter
# notebook -=› display inLine, enabling interactivity. (Awesome.) import matplotlib as mpl
import matplotlib.pyplot as plt
Lil'Log Weng, Lilian. (Jul 2021). What are diffusion models? Lil’Log. https://lilianweng.github.io/posts/2021-07-11-diffusion-models/. SuperAnnotate: introduction to diffusion models for ML
modele dyfuzyjne¶
Wszystkie modele mają problemy, np. GANy są niestabilne w uczeniu.
Modele dyfuzyjne są inspirowane modelami termodynamiki, które nie są w stanie stabilnym (non-equilibrium thermodynamics).
Definiują łańcuch Markowa wiele kroków dyfuzji przez powolne dodawanie szumu do danych, a poźniej uczą się procesu odwrotnego do rekonstrukcji.
Modele są uczone według ustalonej procedury, nie stochastycznie, z przestrzenią ukrytą o wymiarze danych. Proces jest dwu fazowy: w przód (forward) i z porotem (reverse).
Proces forward¶
Forward¶
- Dla $x_0\sim\,q(x)$, proces dyfuzji forward tworzy zaszumione próbki $x_1,\dots,x_T$
- Kroki są kontrolowane schematem wariancji $\{\beta_t\in(0,1)\}_{t=1}^T$
$$q(\mathbf{x}_t \vert \mathbf{x}_{t-1}) = \mathcal{N}(\mathbf{x}_t; \sqrt{1 - \beta_t}\, \mathbf{x}_{t-1}, \beta_t\mathbf{I}) \quad q(\mathbf{x}_{1:T} \vert \mathbf{x}_0) = \prod^T_{t=1} q(\mathbf{x}_t \vert \mathbf{x}_{t-1}),$$ gdzie $\mu_t=\sqrt{1 - \beta_t} \mathbf{x}_{t-1}$ staje się średnią a $\beta_t$ rośnie.
- Próbki coraz bardziej tracą wyraźne cechy
- Proces odwrotny $q(x_{t-1}\mid x_t)$ jest nieznany $$\begin{aligned} \mathbf{x}_t &= \sqrt{1-\beta_t}\mathbf{x}_{t-1} + \sqrt{1 - (1-\beta_t)}\boldsymbol{\epsilon}_{t-1} \end{aligned} $$ niech $\alpha_t=1-\beta_t, \overline\alpha_t=\prod_{t=1}^t\alpha_t, \boldsymbol{\epsilon}_{t-1}, \boldsymbol{\epsilon}_{t-2}, \dots \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ $$ \begin{aligned} \mathbf{x}_t &= \sqrt{\alpha_t}\mathbf{x}_{t-1} + \sqrt{1 - \alpha_t}\boldsymbol{\epsilon}_{t-1} & \\ &= \sqrt{\alpha_t \alpha_{t-1}} \mathbf{x}_{t-2} + \sqrt{1 - \alpha_t \alpha_{t-1}} \bar{\boldsymbol{\epsilon}}_{t-2} & \\ &= \dots \\ &= \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\boldsymbol{\epsilon} \end{aligned} $$ stąd całą zależność $x_t$ od $x_0$ można zapisać $$ \begin{aligned} q(\mathbf{x}_t \vert \mathbf{x}_0) &= \mathcal{N}(\mathbf{x}_t; \sqrt{\bar{\alpha}_t} \mathbf{x}_0, (1 - \bar{\alpha}_t)\mathbf{I}) \end{aligned} $$
kolejne kroki¶
- Łącząc dwa rozkłady normalne z różną wariancją, dostajemy odchylenie $$\mathcal{N}(0,(\sigma_1^2+\sigma_2^2)\mathbb{I})$$ a stąd tutaj $$\sqrt{(1-\alpha_t)+\alpha_t(1-\alpha_{t-1})}=\sqrt{1-\alpha_t\alpha_{t-1}}$$ a po wielu krokach odchylenie $\sqrt{1-\prod_{t=1}^T}=\sqrt{1-\bar{\alpha}_T}$.
- to przypomina modyfikację nauczania SGD gdzie do gradientu dodawany jest jeszcze szum utrudniający utknięcie procesu w płytkim lokalnym minimum.
- w procesie wykorzystywany jest ciąg $\beta_t=1-\alpha_t$,
- zwykle od $\beta_1=10^{-4}$ do $\beta_T=0.02$.
- jest szereg modyfikacji, które poprawiają znacznie zbieżność. Szereg wariancji może być ustalony, może jednak być uczony w trakcie (pod warunkiem wzrostu).
Proces odwrotny (reverse)¶
- Dla bardzo dużych $T$ rozkład coraz bardziej będzie przypominać $\mathcal{N}(0, \mathbb{I})$,
- stąd rozpoczynając proces odwrotny będzie można z tego rozkładu losować
Gausowskość¶
Dla wystarczająco małych $\beta_t$, wartość $q(x_{t-1}\mid x_t)$ także będzie Gausowska.
jednak ta estymacja wymagałaby znajomości całego zbioru danych.
Możemy jednak estymować $$p_\theta(x_{0:T})=p(x_T)\prod_{t=1}^Tp_\theta(x_{t-1}\mid x_t)$$ dla $$p_\theta(x_{t-1}\mid x_t)=\mathcal{N}(x_{t-1}; \mu_\theta(x_t,t),\Sigma_\theta(x_t,t))$$
uczenie polega wiec na znalezieniu takiego $p_\theta()$, który będzie aproksymował proces odwrotny $$ \begin{aligned} q(\mathbf{x}_{t-1} \vert \mathbf{x}_t, \mathbf{x}_0) &= q(\mathbf{x}_t \vert \mathbf{x}_{t-1}, \mathbf{x}_0) \frac{ q(\mathbf{x}_{t-1} \vert \mathbf{x}_0) }{ q(\mathbf{x}_t \vert \mathbf{x}_0) } \\ \end{aligned} $$ i wymagać będzie wyznaczenia odpowiednich $\mu, \epsilon_t$.
Okazuje się, że proces ma uwarunkowania bardzo podobne do VAE,
- stad można użyć dolnego ograniczenia ELBO i odpowiednio przeprowadzić nauczanie.
Koszt kolejnych kroków i optymalizacja¶
- Koszt można przedtawić jako sumę kosztów poszczególnych kroków $$L_\text{VLB} = L_T + L_{T-1} + \dots + L_0,$$ gdzie $$ \begin{aligned} L_T &= KL(q(\mathbf{x}_T \vert \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_T)) \\ L_t &= KL(q(\mathbf{x}_t \vert \mathbf{x}_{t+1}, \mathbf{x}_0) \parallel p_\theta(\mathbf{x}_t \vert\mathbf{x}_{t+1})) \text{ for }1 \leq t \leq T-1 \\ L_0 &= - \log p_\theta(\mathbf{x}_0 \vert \mathbf{x}_1). \end{aligned} $$
- Wszystkie elementy dywergencji są przedstawione w postaci zamkniętej,
- nie wymagają samplowania dla porównania rozkładów.
- Dla $L_0$ potrzebny jest dodatkowy dekoder.
- Poszczególne elementy $L_t$ są możliwe do sparametryzowania przez $\mu_t$ oraz $\Sigma_\theta$ i ich predykcję.
- Okazuje się, ze nawet uproszczony model tej predykcji działa lepiej i stabilniej.
Inne procesy odszumiania¶
- Ten proces jest podobny do procesów odszumiania, np. Denoising Auto Encoder DAE.
- W DAE do przykładów jest dodawany szum, by wymusić jego redukcję, a przez to nauczenie się prawdziwego rozkładu danych.
Warunkowy proces generowania¶
uczenie klasyfikatora $f_\theta(y\mid x_t,t)$ na zaszumionym obrazie
wykorzystuje gradient $$\nabla_x\log\,f_\theta(y\mid x_t)$$ sterujący procesem samplowania w kierunku etykiety warunkującej $y$.
Proces jest wzmocnieniem samego samplingu.
- Algorytm ablated diffusion model ADM osiągnął lepsze wyniki niż najlepszy duży model BigGAN.
Bez explicite nauczonego klasyfikatora też jest możliwe uzyskanie warunkowej generacji. Możemy mieć
- bezwarunkowy proces odszumiający $p\theta(x)$ zbudowany przez estymator zysku (score) $\epsilon_\theta(x_t,t)$ oraz
- model warunkowy $p\theta(x\mid y)$.
- Te modele mogą być nauczone przez jedną sieć neuronową tak, że $p_\theta(x\mid y)$ jest uczony na parach $(x, y)$.
- Jednak wartość warunkowa $y$ jest losowo usuwana tak, ze model będzie w stanie generować obiekty także bezwarunkowo.
- Taki model osiąga dobry balans między FID (dobre cechy i rozróżnienie miedzy obiektami prawdziwymi i generowanymi) oraz IS (zróżnicowanie i jakość).
- Model GLIDE wykorzystywał model CLIP jako klasyfikator oraz sterowanie bez klasyfikatora znajdując drugie podejście jako zwykle lepsze.
Przyspieszenie procesu¶
- Cały proces tworzenia modelu może być bardzo wolny, gdyż liczba kroków $T$ może być bardzo duża.
- Można zmniejszyć liczbę kroków, poprawiając model co $n$ oryginalnych kroków.
50 tysięcy obrazów 32x32 może zająć 20 godzin dla procesu Markowa, ale mniej niż minutę dla modelu GAN.
Destylacja¶
- W procesie uczenia modelu (w przód, forward) można wykorzystać destylację
- proces uczy się kolejnych kroków $q(x_t\mid x_{t-1}$
- po nauczeniu dwóch kolejnych następuje destylacja
- dwa kolejne służą jako nauczyciel
- uczą jeden kolejny proces, który bedzie przyspieszał pracę
Skalowanie rozdzielczości oraz jakości¶
- Możliwością jest stworzenie strumienia (pipeline) przez utworzenie szeregu modeli dyfuzji w coraz wyższych rozdzielczościach.
- Podstawowym elementem jest augmentacja oparta szumem pomiędzy elementami.
- Będzie to polegać na dodaniu silnej augmentacji danych na warunkującym wejściu $z$ dla każdego $p_\theta(x\mid z)$.
- bardzo częstym wyborem jest U-net (o tym poniżej). ![[cascaded-diffusion.png]] [Rysunek Ho] Najlepiej dodać szum Gausowski dla niskich rozdzielczości , oraz Gaussian blur w wysokich (zamiast uśrednionego, daje większą wagę centralnym pikselom; inaczej Gausowskie wygładzanie).
- Dla obrazu niskiej rozdzielczości proces jest zatrzymywany dla niskiego $t$.
- Odwrotny proces niskiej rozdzielczości dochodzi do punktu $0$
- Wtedy jest zaburzany przez $z_t\sim q(x_t\mid x_0)$, a ten podaje $z_t$ do modelu wyższej rozdzielczości.
Model unCLIP¶
Przy generowaniu obrazów sterowanych tekstem wydajne jest użycie CLIP.
Dla CLIP uczonego parami $(x, y)$, gdzie $x$ to obraz, a $y$ to związana etykieta,
- budujemy embeddingi tekstu $c^t(y)$ i obrazu $c^i(x)$. unCLIP uczy się równolegle modeli
- model prior $P(c^i\mid y)$ dla zadanego $y$ zwraca embedding (dla CLIP) $c^i$ tekstu $y$,
- dekoder $P(x\mid x^i, [y])$ generuje obraz $x$ mając embedding $c^i$; ewentualnie oryginalny tekst $y$.
Modele pozwalają na generowanie warunkowe $$P(x\mid y) = P(x, c^i\mid y) = P(x\mid c^i,y)\,P(c^i\mid y).$$
[Ramesh 2022]
Niektóre inne modele wykorzystują LLM dla kodowaniaopisu obrazu przy jego generowaniu.
- Większa waga dla dobrego dopasowania tekstu prowadzi często do słabszej jakości obrazu. Konieczny jest często tuning.
Architektura modeli¶
Podstawowymi dla modeli dyfuzyjnych są U-Net i Transformer.
- U-Net
[Ronneberger]
- Transformer
Transformer¶
Tu tworzony jest Diffusion Transformer (DiT), który
- jako wejście bierze reprezentację typu latent,
- ta jest dzielona na $p^2$ kawałków (patchy), jak w ViT,
- sekwencja tych patchy jest podawana na wejście transformera,
- między embeddingiem a wartwą normalizacji warstwowej (LayerNorm) i wartwą liniową znajduje się dodatkowy blok DiT,
- blok składa przypomina blok Transformera i przewiduje parametry potrzebne do wygenerowania szumu kolejnej transformacji dyfuzji,
- bloków Transformera, z dodanym DiT, jest N. Model Transformera jest wyjątkowo elastyczny, jednak w połączeniu z koniecznym długim łańcuchem Markowa w przetwarzaniu dyfuzji, może być kosztowny.
Przykładowe uczenie (denoising-diffusion-pytorch)¶
import torch
from denoising_diffusion_pytorch import Unet, GaussianDiffusion
/Users/igor/miniforge3/envs/mini/lib/python3.10/site-packages/torchvision/io/image.py:14: UserWarning: Failed to load image Python extension: 'dlopen(/Users/igor/miniforge3/envs/mini/lib/python3.10/site-packages/torchvision/image.so, 0x0006): Library not loaded: @rpath/libjpeg.9.dylib Referenced from: <0B7EB158-53DC-3403-8A49-22178CAB4612> /Users/igor/miniforge3/envs/mini/lib/python3.10/site-packages/torchvision/image.so Reason: tried: '/Users/igor/miniforge3/envs/mini/lib/python3.10/site-packages/torchvision/../../../libjpeg.9.dylib' (no such file), '/Users/igor/miniforge3/envs/mini/lib/python3.10/site-packages/torchvision/../../../libjpeg.9.dylib' (no such file), '/Users/igor/miniforge3/envs/mini/lib/python3.10/lib-dynload/../../libjpeg.9.dylib' (no such file), '/Users/igor/miniforge3/envs/mini/bin/../lib/libjpeg.9.dylib' (no such file)'If you don't plan on using image functionality from `torchvision.io`, you can ignore this warning. Otherwise, there might be something wrong with your environment. Did you have `libjpeg` or `libpng` installed before building `torchvision` from source? warn(
# Definicja modelu Unet
# info wymiar map cech oraz mnożniki dla kolejnych kompresji
model = Unet(
dim = 64,
dim_mults = (1, 2, 4, 8)
)
# Model dyfuzji
diffusion = GaussianDiffusion(
model,
image_size = 128,
timesteps = 1000, # number of steps
# loss_type = 'l1' # L1 or L2
)
# uczenie z szumu
training_images = torch.randn(8, 3, 128, 128)
loss = diffusion(training_images)
loss.backward()
# generowanie obrazów
sampled_images = diffusion.sample(batch_size = 4)
sampling loop time step: 0%| | 0/1000 [00:00<?, ?it/s]
import torchvision
to_pil = torchvision.transforms.ToPILImage()
def show(img):
npimg = img.numpy()
plt.imshow(np.transpose(npimg, (1, 2, 0)), interpolation='nearest')
for k in range(4):
show(sampled_images[k])
# uczenie dla zbioru danych
from denoising_diffusion_pytorch import Unet, GaussianDiffusion
from denoising_diffusion_pytorch import Trainer
model = Unet(
dim = 64,
dim_mults = (1, 2, 4, 8)
)
diffusion = GaussianDiffusion(
model,
image_size = 128,
timesteps = 1000, # number of steps
# loss_type = 'l1' # L1 or L2
)
trainer = Trainer(
diffusion,
'/Users/igor/tmp',
train_batch_size = 32,
train_lr = 2e-5,
train_num_steps = 700000, # total training steps
gradient_accumulate_every = 2, # gradient accumulation steps
ema_decay = 0.995, # exponential moving average decay
amp = True # turn on mixed precision
)
--------------------------------------------------------------------------- ValueError Traceback (most recent call last) Cell In[38], line 1 ----> 1 trainer = Trainer( 2 diffusion, 3 '/Users/igor/tmp', 4 train_batch_size = 32, 5 train_lr = 2e-5, 6 train_num_steps = 700000, # total training steps 7 gradient_accumulate_every = 2, # gradient accumulation steps 8 ema_decay = 0.995, # exponential moving average decay 9 amp = True # turn on mixed precision 10 ) File ~/miniforge3/envs/mini/lib/python3.10/site-packages/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py:910, in Trainer.__init__(self, diffusion_model, folder, train_batch_size, gradient_accumulate_every, augment_horizontal_flip, train_lr, train_num_steps, ema_update_every, ema_decay, adam_betas, save_and_sample_every, num_samples, results_folder, amp, mixed_precision_type, split_batches, convert_image_to, calculate_fid, inception_block_idx, max_grad_norm, num_fid_samples, save_best_and_latest_only) 906 super().__init__() 908 # accelerator --> 910 self.accelerator = Accelerator( 911 split_batches = split_batches, 912 mixed_precision = mixed_precision_type if amp else 'no' 913 ) 915 # model 917 self.model = diffusion_model File ~/miniforge3/envs/mini/lib/python3.10/site-packages/accelerate/accelerator.py:485, in Accelerator.__init__(self, device_placement, split_batches, mixed_precision, gradient_accumulation_steps, cpu, dataloader_config, deepspeed_plugin, fsdp_plugin, megatron_lm_plugin, rng_types, log_with, project_dir, project_config, gradient_accumulation_plugin, step_scheduler_with_optimizer, kwargs_handlers, dynamo_backend, deepspeed_plugins) 481 self.native_amp = True 482 if self.device.type not in ("xpu", "cuda", "npu", "xla", "mlu", "musa") or is_torch_xla_available( 483 check_is_tpu=True 484 ): --> 485 raise ValueError(f"fp16 mixed precision requires a GPU (not {self.device.type!r}).") 486 kwargs = self.scaler_handler.to_kwargs() if self.scaler_handler is not None else {} 487 self.scaler = get_grad_scaler(self.distributed_type, **kwargs) ValueError: fp16 mixed precision requires a GPU (not 'mps').
trainer.train()
--------------------------------------------------------------------------- NameError Traceback (most recent call last) Cell In[39], line 1 ----> 1 trainer.train() NameError: name 'trainer' is not defined